require 'config'
require 'lib'

local rulematch = ngx.re.find
local USER_AGENT_RULES = get_rule('useragent.rule')
local SERVER_NAME_LIMIT = read_serlimit_config()

function conn_limit()
    local limit_conn = require "resty.limit.conn"

    local http_host = ngx.var.http_host
    ser_conn_limit = SERVER_NAME_LIMIT["ConnLimit"][http_host]
    if ser_conn_limit ~= nil then
        maxconn,burst,reqtime = string.match(ser_conn_limit,'(.*)/(.*)/(.*)')
    else
	maxconn,burst,reqtime = string.match(GlobalConnLimit,'(.*)/(.*)/(.*)')
    end
    if maxconn == nil or burst == nil or reqtime == nil then
        maxconn,burst,reqtime = string.match(GlobalConnLimit,'(.*)/(.*)/(.*)')
    end
    local lim, err = limit_conn.new("my_limit_conn_store", tonumber(maxconn), 
    		tonumber(burst), tonumber(reqtime))
    if not lim then
        ngx.log(ngx.ERR,
                "failed to instantiate a resty.limit.conn object: ", err)
	return ngx.exit(500)
    end
    
    local http_host = ngx.var.http_host
    local ATTACK_URI = ngx.var.uri
    local key = http_host..get_client_ip()..ATTACK_URI
    local delay, err = lim:incoming(key, true)
    if not delay then
        if err == "rejected" then
            return say_html("conn")
        end
        ngx.log(ngx.ERR, "failed to limit req: ", err)
        return say_html("conn")
    end
    
    if lim:is_committed() then
        local ctx = ngx.ctx
        ctx.limit_conn = lim
        ctx.limit_conn_key = key
        ctx.limit_conn_delay = delay
    end
    
    local conn = err
    
    if delay >= 0.001 then
        ngx.sleep(delay)
    end

end


function count_limit()
    local limit_count = require "resty.limit.count"

    -- rate: CCcount requests per CCseconds
    local http_host = ngx.var.http_host
    ser_count_limit = SERVER_NAME_LIMIT["CountLimit"][http_host]
    if ser_count_limit ~= nil then
        count,sec = string.match(ser_count_limit,'(.*)/(.*)')
    else
        count,sec = string.match(GlobalCountLimit,'(.*)/(.*)')
    end
    if count == nil or sec == nil then
        count,sec = string.match(GlobalCountLimit,'(.*)/(.*)')
    end

    ngx.log(ngx.ERR, count)
    ngx.log(ngx.ERR, sec)
    local lim, err = limit_count.new("my_limit_count_store", tonumber(count), tonumber(sec))
    if not lim then
        ngx.log(ngx.ERR, "failed to instantiate a resty.limit.count object: ", err)
	return ngx.exit(500)
    end
    local http_host = ngx.var.http_host
    local ATTACK_URI = ngx.var.uri
    local key = http_host..get_client_ip()..ATTACK_URI
    local delay, err = lim:incoming(key, true)

    if not delay then
        if err == "rejected" then
            ngx.header["X-RateLimit-Limit"] = "5000"
            ngx.header["X-RateLimit-Remaining"] = 0
            say_html("count")
        end
        ngx.log(ngx.ERR, "failed to limit count: ", err)
        say_html("count")
    end

    local remaining = err

    ngx.header["X-RateLimit-Limit"] = "5000"
    ngx.header["X-RateLimit-Remaining"] = remaining
end

--deny user agent
function user_agent_attack_check()
    if config_user_agent_check == "on" then
        local USER_AGENT = ngx.var.http_user_agent
        if USER_AGENT ~= nil then
            for _,rule in pairs(USER_AGENT_RULES) do
                if rule ~="" and rulematch(USER_AGENT,rule,"jo") then
                    if config_waf_enable == "on" then
                        waf_output()
                        return true
                    end
                end
            end
        end
    end
    return false
end
